import cv2
import numpy as np


class CalculateDistance:
    """计算滑块验证码需要滑动的距离"""

    def __init__(self, background_path, slide_path, offset_px, offset_py, display):
        """
        初始化
        :param background_path: 验证码背景大图路径
        :param slide_path: 滑块小图路径
        :param offset_px: 小图距离大图左侧的像素偏移量
        :param offset_py: 小图距离大图顶部的像素偏移量
        """
        # 读取图片
        self.backgound_img = cv2.imread(background_path)
        self.offset_px = offset_px
        self.offset_py = offset_py
        self.slide_img = cv2.imread(slide_path, cv2.IMREAD_UNCHANGED)

        # 计算x轴缩放因子，以50px为基准
        scale_x = 50 / self.slide_img.shape[1]

        # 使用最近邻插值法缩放，得到缩放后的50x50的图片
        self.slide_scale_img = cv2.resize(self.slide_img, (0, 0), fx=scale_x, fy=scale_x)

        self.background_cut_img = None
        self.display = display

    def get_distance(self):
        """计算滑动距离"""
        # 将小图转为灰度
        slide_grey_img = cv2.cvtColor(self.slide_scale_img, cv2.COLOR_BGR2GRAY)

        # 使用Canny算子，提取图片边缘特征
        slide_edge_img = cv2.Canny(slide_grey_img, 100, 250)
        # self.cv_show('canny', slide_edge_img)

        # 将背景图转换为灰度
        background_grey_img = cv2.cvtColor(self.background_cut_img, cv2.COLOR_BGR2GRAY)

        # 使用Canny算子提取背景图的边缘特征
        background_edge_img = cv2.Canny(background_grey_img, 200, 300)
        # self.cv_show('bg_canny', background_edge_img)

        # 获取滑块图像的尺寸
        h, w = slide_edge_img.shape

        # 将滑块图与背景进行模板匹配，找到缺口所在的位置
        result = cv2.matchTemplate(background_edge_img, slide_edge_img, cv2.TM_CCOEFF_NORMED)

        # 获取匹配结果中的最佳位置
        min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)

        # 获取缺口左上角的位置
        top_left = (max_loc[0], max_loc[1])

        # 右下角位置
        bottom_right = (top_left[0] + w, top_left[1] + h)

        # 调试模式下显示匹配结果
        if self.display:
            print(top_left)
            print(bottom_right)
            after_img = cv2.rectangle(self.background_cut_img, top_left, bottom_right, (0, 0, 255), 2)
            self.cv_show('after', after_img)

        # 计算滑动距离
        slide_distance = top_left[0] + w + 10  # 加上10像素作为偏移
        return slide_distance

    def cut_background(self):
        """裁剪背景图中无关区域"""
        height = self.slide_scale_img.shape[0]

        # 将背景图中上下多余部分以及滑块图片部分去除
        self.background_cut_img = self.backgound_img[self.offset_py - 10:self.offset_py + height + 10,
                                  self.offset_px + height + 10:]

    def cv_show(self, name, img):
        """显示图像（调试用）"""
        cv2.imshow(name, img)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

    def run(self):
        """执行计算流程"""
        self.cut_background()
        return self.get_distance()


if __name__ == '__main__':
    # 示例用法
    path = r'D:\Python Files\Spider\Pictures\sildes\\'
    background_path = path + r'BG-1663996834.jpg'
    slide_path = path + r'SD-1663996834.jpg'
    offset_px = 28  # 小图左侧偏移
    offset_py = 28  # 小图顶部偏移

    # 实例化并计算距离
    main = CalculateDistance(background_path, slide_path, offset_px, offset_py, display=True)
    distance = main.run()
    print(f"需要滑动的总距离: {distance} 像素")
